#!/usr/bin/env python
""" 
Basile Clout, September 2nd, 2007
wmap is the client mapping the grid computing
Work in coordination with the pywrekamapd daemon, which must be present in the path on the machines in the considered network

Command:
-h, --help Display this help
-c arg, --conf arg Configuration file (default: wmap.conf)

Configuration options
Several parameters can be configured in a configuration file.
Default values are hard coded, and correspond to the lids (lids.cs.unb.ca) configuration
Each line starting with a "\#" is discarded (comments)
Each line (not starting with the # character) must correspond to the format:
keyword = value
The following keywords are useful (you can declare other keywords if you want)
wmapd_bin	Path to the wmap daemon binary (pywrekamapd) on the clients (DEFAULT: /root/config/code/tools/pywrekamapd.py)
wrekaUDP_bin	Path to the wrekaUDP daemon on the client (DEFAULT: /root/config/code/tools/wrekaUDP.pl)
ssh_bin		Path to the ssh binary (DEFAULT: /usr/bin/ssh)
iperf_bin	Path to the iperf bin (DEFAULT: /usr/local/bin/iperf)
ping_bin	Path to the ping bin (DEFAULT: /bin/ping)
log_file	Path and name of the log file (DEFAULT: wmap.log)
pagrid_out	Path and name of the pagrid output file (DEFAULT: wmap.out)
mpirun_bin	Path to the mpirun binary (DEFAULT: /usr/local/bin/mpirun)
wrekaMPI_bin	Path to the wrekaMPI binary (MPI bench), used by wmapd and wmap (DEFAULT: /home/basil/mpi/netMPI/wreka_MPI)
tmp_dir		Path to the temporary directory (have to get read and write rights) (DEFAULT: /tmp/wmap/
hostlist	List of hosts in the vitual grid (DEFAULT: lids02 lids04 lids06 lids09 lids13 lids15 lids16)
hosts_file	Name of the tmp file containing the host names, in a format adapted to the mpirun --hostfile program. Built from hostlist. Located in tmp_dir (DEFAULT: conf[\"tmp_dir\"] + \"hostfile.txt\")
ratio		Second line of the pagrid grid input file (pagrid_out) (DEFAULT: 0 1 0.25)
cpu_user	User under which the euid the CPU test in wmapd is being ran. Correspond to the user name given in the wrekavoc configuration file (DEFAULT: basil)
flags_wmapd	Parameters to pass to the daemon wmapd. DO NOT TOUCH! (DEFAULT: --net_off --set_user conf[\"cpu_user\"]


For bugs reports, comments, suggestions, donations :D, ... please contact me at
basile.clout@unb.ca
"""

import os
import string
import pprint
import time
import socket
from datetime import datetime
import pwd
import sys
import getopt

class node:
    """ Description of a node 
    """
    
    def __init__(self, rank):
        """ Initialize a node. Compatible with the definitions of a node in the daemon
        """
        self.name = ""
        self.address = ""
        self.neighbors = []
        self.cpu = 0
        self.user = ""
        self.rank = rank
        self.links_tasks ={}
        self.links_config = {}
        self.links_results = {}
        
    def getRank(self, address, log):
        """ Get rank of one node, given its address or its name 
        For this node.
        """
        if len(self.links_tasks) == 0:
            log.write("ERROR: struct links_tasks empty\n")
        elif (address == ""):
            log.write("ERROR: Please give an address or a name to getRank\n")
        else:

            try:
                for rank, names in self.links_tasks.iteritems():
                    if address in names:
                        return int(rank)
            except:
                log.write("ERROR: Impossible to walk the links_tasks struct in getRank\n")
                
            log.write("Warning: impossible to get the rank\n")
            return "none"


class grid:
    """ Description of a grid 
:    """
    
    def __init__(self, conf, log):
        self.hostlist = conf["hostlist"].split()            # Here list is a real list
        self.nbprocs = len(self.hostlist)
        self.nblinks = 0
        self.links = []                            # Links: contains all the links in the cluster; (node a, node b)
        self.conf = conf
        self.log = log
        self.nodes = {}
        self.master_node = conf["master_node"]   # Index (entry) corresponding to the master node in the grid
        self.hostfile = self.buildHostfile()
        self.ratio = conf["ratio"].split()
        
        
        
    def getRank(self, address):
        """Get the rank of a node known by its address or hostname
        (NOT very good: address has to correspond to the format of hostlist)
        Use the getRank of the node class instead
        """
        print "address = _" + address +"_"
        print "hostlist = _" + str(self.hostlist) + "_"
        if (address != ""):
            try:
                return self.hostlist.index(address)
            except:
                self.log.write("ERROR: No such "+ address + " name in the hostlist " + "\n")
                return "none"
        else:
            self.log.write("ERROR: Please specify an address or a hostname" + "\n")
    

    def getInfoNodes(self):
        
        print("Get Info from the nodes in " + string.join(self.hostlist))
        self.log.write("Get Info from the nodes in " + string.join(self.hostlist) + "\n")
        
        order_node = self.conf["wmapd_bin"] + " " + self.conf["flags_wmapd"]
        t = string.Template(self.conf["ssh_bin"] + " $address \"" + order_node + "\" ")
        
        tMPI = string.Template(self.conf["mpirun_bin"] \
                                       + " -np " + str(self.nbprocs) + \
                                       " --hostfile " + self.conf["hosts_file"] + \
                                       " " + self.conf["wrekaMPI_bin"] + " -m $master")
                
        i = 0
        for name in self.hostlist:
            
            try:
                _node = node(self.hostlist.index(name))
                _node.name = name
            except:
                print "ERROR: " + name + " is not in the host list!\n"
                self.log.write("ERROR: " + name + " is not in the host list!\n") 
            
            fl_hasnet = self.getLocalInfo(_node, t, name)
            if fl_hasnet != 2:
                print "Get local info for "+ name + " OK"
                self.log.write("Get Local info for " + name + " OK!\n")
            else:
                print "Error getting local info"
                self.log.write("Error getting local info\n")
                sys.exit(1)
                
            if (fl_hasnet == 0):
                r = self.getLinksMpiInfo(_node, tMPI, name)
                if r == 0:
                    print "Get Links info (using MPI) OK!"
                    self.log.write("Get Links info (using MPI) OK!\n")
                else:
                    print "Error getting links infos"
                    self.log.write("Error getting links infos")
            else:
                print "ERROR: case not yet handled!"
                self.log.write("ERROR: case not yet handled!\n")
                    
            self.nodes[_node.rank] = _node
            
        self.nblinks = self.getLinksNumber()


    def getLocalInfo(self, _node, t, name):
        """ Get local info corresponding to the node _node 
        - cpu
        - neighbors
        """
        
        order = t.substitute(address = name)
        print "Request info from %s: \t%s ..." % (name, order)
        self.log.write("Info from %s: \t%s ...\n" % (name, order))

        stdin, stdout, stderr = os.popen3(order)
        output, errput = stdout.read(), stderr.read()
            
        if errput == "":
            print "OK!"
            self.log.writelines(output + "\n OK!\n")
        else:
            print "ERROR: ssh order to %s failed. See log file" % name
            self.log.write(output + "\n")
            self.log.write("ERROR: ssh order to " + name + " failed: \n")
            self.log.writelines(errput + "\n")
            sys.exit(1)
          
        fl_hasnet = self.parse_wmapd(_node, output, self.log)

        return fl_hasnet
    
            
    def getLinksMpiInfo(self, _node, tMPI, name):
        """ Get Links info concerning the node _node
        by launching wrekaMPI on the master node and 
        adapting the master node to be _node
        """
        
        #master_node = self.hostlist.index("name") 
        order = tMPI.substitute(master = _node.rank)
        print "Get links info for %s: \t %s" % (name, order)
        self.log.write("Launch wrekamap on " + name + " with :\n\t" + order + "\n")
        
        stdin, stdout, stderr = os.popen3(order)
        output, errput = stdout.read(), stderr.read()

        if errput != "":
            print "ERROR: launching wrekaMPI"
            self.log.write("ERROR when launching wreka_MPI!:\n")
            self.log.writelines(errput)
            sys.exit(1)
        else:
            self.log.write(" ... wreka_MPI launch OK!\n")
        
        self.log.write("Parse wreka_MPI output ...")
        
        if (self.parse_wrekaMPI(_node, output, self.log) == 1):
            self.log.write("Error when parsing wrekaMPI output!\n")
        else:
            print "OK!"
            self.log.write("OK!\n")
            self.log.write("links_config list:\n")
            pprint.pprint(_node.links_config, self.log)
            self.log.write("links_tasks dictionary:\n")
            pprint.pprint(_node.links_tasks, self.log)
            self.log.write("links_results dictionary:\n")
            pprint.pprint(_node.links_results, self.log)

        return 0

                
    def buildHostfile(self):
        """ Build a hostfile suitable for use by the MPI implementation
        from the string list containing all the hosts (or ip addresses)
        into the file \"hosts_file\" 
        """
        self.log.write("Building hostfile ...")
        hostfile = file(self.conf["hosts_file"], "w")
        for line in self.hostlist:
            hostfile.writelines(line + "\n")
        self.log.writelines("hostfile \"" + self.conf["hosts_file"] + "\" written OK!\n") 
                
        return self.conf["hosts_file"]                                                           
    
    
    def parse_wrekaMPI(self, _node, output, log):    
        """ Parse the output of wrekaMPI, 
        and fills up the data structures of the node object 
        """
        
        next = 0 
        count = 0
        nbtasks = 0
        
        for line in output.splitlines():
            
            parts = line.split()
            
            if next == 0 and parts[0] == "config":
                next = 1                        # para "config"
            elif next == 0 and parts[0] == "tasks":
                next = 2                        # para "tasks"
            elif next == 0 and parts[0] == "results":
                next = 3                        # para "results"
            elif next == 0:
                log.write("Error during parsing. Pb with the paragraph's order\n")
                return 1
            else:                                
                """ We are inside a paragraph """
                if next == 1:                    # in para "config"
                    if len(parts) != 11:
                        log.write("Error when parsing the \"config\" para\n")
                        return 1
                    else:
                        _node.links_config["master"]=parts[0]
                        _node.links_config["nbtasks"]=parts[1]
                        _node.links_config["method"]=parts[2]
                        _node.links_config["size"]=parts[3]
                        _node.links_config["latsize"]=parts[4]
                        _node.links_config["nbtests"]=parts[5]
                        _node.links_config["skip"]=parts[6]
                        _node.links_config["isrrequests"]=parts[7]
                        _node.links_config["help"]=parts[8]
                        _node.links_config["debug"]=parts[9]
                        _node.links_config["live"]=parts[10]
                        
                        if _node.links_config["help"]==1:
                            log.write("Warning: Help mode enabled!\n")
                        if _node.links_config["debug"]==1:
                            log.write("Warning: Debug mode enabled!\n")
                        if _node.links_config["live"]==1:
                            log.write("Info: Live mode enabled!\n")
                        
                        nbtasks = int(_node.links_config["nbtasks"])
                        next = 0
                        
                elif next == 2:
                    if nbtasks == 0:
                        log.write("Error in paragraph order (tasks)\n")
                        return 1
                    if len(parts) != 3:
                        log.write("Error when parsing the \"tasks\" para\n")
                        return 1
                    
                    _node.links_tasks[parts[0]]=parts[1:]        # The key is the address (from _node.neighbors)
                    count+=1
                    
                    if count == nbtasks:
                        next = 0
                        count = 0
                
                elif next == 3:
                    if nbtasks == 0:
                        log.write("Error in paragraph order (results)")
                        return 1
                    if len(parts) != 10:
                        log.write("Error when parsing the \"results\" para\n")
                        return 1
                    
                    _node.links_results[parts[1]]=[float(x) for x in parts[2:]]        # the key is the rank (from _node.link_tasks)
                    count+=1
                    
                    if count == nbtasks:
                       next = 0
                       count = 0 
                       
        return 0
    
    
    
    def parse_wmapd(self, _node, output, log):
        """ 
        Parse the output of wrekamapd, following the indications of the 
        flags on the first line of the output
        """
        print output
        flags = []
        for line in output.splitlines():

            parts = line.split()

            if parts[0] == "node" and flags == []:                   
                if _node.name not in parts:
                    log.write("Error name of the node: " + parts[1]+"/" + parts[2] + " vs " + node.name + "\n")
                _node.address = parts[2]
                flags = parts[3:]
                log.write("Parse node line ok in wmapd output\n")
                   
                if flags[3] == '1':
                    log.write("Error in wmapd.\n")
                    print "Error in wmapd"
                    sys.exit(1)
 
            elif flags[0] == '1' and parts[0] == "cpu":
                _node.cpu = float(parts[1])
                _node.user = parts[2]
                log.write("Parse cpu line ok in wmapd output\n")
                
            elif flags[1] == '1' and parts[0] == "neighbors":
                _node.neighbors = parts[1:]
                
            elif flags[2] == '1' and parts[0] == "to":
                log.write(line + "\n")
                print "Case not yet implemented\n"
                log.write("ERROR: case [Network in wmapd] not yet implemented\n")
                
            else:
                log.write("Error when parsing output from wmapd\n")
                print "ERROR: when parsing wmapd output"
                
                
        return int(flags[3])
    
    
    def getLinksNumber(self):
        """ Returns the number of links available in the virtual grid
        """
        nblinks = 0
        links = []
        for _node in self.nodes.itervalues():
            #            print _node.address, _node.name
            for address in _node.neighbors:
                rank = _node.getRank(address, self.log)
                if rank == "none":
                    continue
                else:           # The node at rank "rank" is a on a real link
                    if (_node.rank, rank) not in links:
#                        print str((_node.rank, rank)), str(links)
                        links.append((_node.rank, rank))
                        links.append((rank, _node.rank))
                        nblinks += 1
        
        self.log.write("Number of links in the grid: " + str(nblinks) + "\n")
        return nblinks
     

    def printSummary(self);
    """ Format and return a nice representation of the network
    """
     
class pagrid:
    """ This class defines functions useful for formatting the input for pagridL
    """       
    
    def __init__(self, grid, log):
        self.grid = grid
        self.log = log
        self.mincpu = self.getMinCPU()
        self.maxbdw = self.getMaxBdw()
        self.master = grid.master_node
        
        
    def getMinCPU(self):
        """ Returns the minimum CPU power 
        which is in fact the maximum CPU time!
        for the nodes in the grid.
        Used to compute the relative weight of the node
        """
        
        minCPU = self.grid.nodes[0].cpu
        for _node in self.grid.nodes.itervalues():
            if _node.cpu > minCPU:
                minCPU = _node.cpu

        self.log.write("min cpu = " + str(minCPU) + "\n")
        return minCPU
    
    def getMinNet(self):
        """ Returns the minimum latency time (in ms)
        and minimum bandwidth (in MBits/s)
        for all the links in the grid.
        Used to compute the relative weight of the links
        """
        
        minLat = 10000
        minBdw = 10000
        for _node in self.grid.nodes.itervalues():
            for res in _node.links_results.itervalues(): 
#                print "res " + str(res[0]), str(res[4]) 
                if res[0] < minLat:
                    minLat = res[0]
                if res[4] <minBdw:
                    minBdw = res[4]
                
        self.log.write("min latency = " + str(minLat) + "\n")
        self.log.write("min bandwidth = " + str(minBdw) + "\n")
        return minLat, minBdw
    

    def getMaxBdw(self):
        """ Returns the maximum bandwidth in the cluster
        """

        maxBdw = 0
        for _node in self.grid.nodes.itervalues():
            for res in _node.links_results.itervalues():
                if res[4] > maxBdw:
                    maxBdw = res[4]

        self.log.write("max bandwidth = "+ str(maxBdw) + "\n")
        return maxBdw

        
    def getWeightCPU(self, cpu):
        """ Returns the cpu weight of the processor
        """
        cpuWeight = self.mincpu/cpu
        return round(cpuWeight)

    
    def getWeightBdw(self, bdw):
        """ Returns the bdw weight of the link
        """
        bdwWeight = self.maxbdw/bdw
        return round(bdwWeight)

    
    def getWeightLat(self, lat):
        """ Returns the latency weight of the link
        """
        latWeight = lat*self.maxbdw*1000
        return round(latWeight)
    
    
    def printPagrid(self, out):
        """ Print an output compatible with PaGridL
            number_of_processors number_of_links
            proc_1 proc_2 ratio
            cpu_weight rank1 bandwidth_weight1 latency_weight1 rank2 bdw_weight2 ...  \tEach line corresponding to one processor, in rank order\n
        """
        print "Print pagrid output started ..."

        weightBdw = 0
        weightLat = 0
        intro = ""
        self.log.write("PaGrid output: \n")
        intro += "%Grid description file for PaGridL\n"
        intro += "%Automatically generated with " + sys.argv[0] + " on " + str(datetime.now()) + " by " + pwd.getpwuid(os.geteuid())[0] + "\n"
        intro += "%Syntax:\n"
        intro += "%#_of_processors #_of_links\n"
        intro += "%proc_1 proc_2 ratio\n"
        intro += "%cpu_weight rank1 bandwidth_weight1 latency_weight1 rank2 bdw_weight2 ...  \tEach line corresponding to one processor, in rank order\n"
        
        out.writelines(intro)
        self.log.writelines(intro)
        
        out.write("%d %d\n" % (self.grid.nbprocs, self.grid.nblinks))
        self.log.write("%d %d\n" % (self.grid.nbprocs, self.grid.nblinks))
        
        if len(self.grid.ratio) == 3:
            #self.log.write("%d %d %d\n" % tuple([int(x) for x in self.grid.ratio]))
            #out.write("%d %d %d\n" % tuple([int(x) for x in self.grid.ratio]))
            self.log.write("%s %s %s\n" % tuple(self.grid.ratio))
            out.write("%s %s %s\n" % tuple(self.grid.ratio))
        else:
            self.log.write("ERROR: Format of ratio: "+ self.grid.ratio + "\n")
        
        for _node in self.grid.nodes.itervalues():
            string = str(int(self.getWeightCPU(_node.cpu))) + " "
            for neigh in _node.neighbors:
                rank = _node.getRank(neigh, self.log) # Use node.links_tasks to get the rank from the value in the neighbors structure
                if rank == "none":
                    self.log.write(neigh + " is a neighbor but is not in the hostlist\n")
                    continue
                res = _node.links_results[str(rank)]
                weightBdw = self.getWeightBdw(float(res[4]))
                weightLat = self.getWeightLat(float(res[0]))
                string += str(rank) + " " + str(int(weightBdw)) + " " + str(int(weightLat)) + " "
            out.write(string + "\n")
            self.log.write(_node.address + " " + string + "\n")
                           
        print "OK!"
        out.close()

    def printGraph(self):
        """ Generate the .dot file
        """
        
        print "Graph generation started ..."
        order = self.grid.conf["wmap2dot"] + " " + self.grid.conf["pagrid_out"] + " > " + self.grid.conf["pagrid_dot"]
        self.log.write("Graph generation started with order: %s\n", order)
        os.system(order)

        print "OK!"
        
       
def parse_conf(handle, conf):
   """ 
   Parse the configuration file and update conf
   """ 
   for line in handle.readlines():
       parts = line.split("=")
       if parts[0][0] == "#":
           continue
       elif len(parts) != 2:
           print "Error when parsing configuration file"
           sys.exit(1)
       else:
           #if len(parts[1].split()) == 1:
           #    conf[parts[0].strip()] = parts[1].strip()
           #else:
                #conf[parts[0].strip()] = parts[1].strip().split()
           conf[parts[0].strip()] = parts[1].strip()

   return conf
            
            
def default_conf(conf, key=""):
    """
     Default configuration (OK for lids.cs.unb.ca)
     """
    def_conf = {}
    def_conf["wmapd_bin"] = "/root/config/code/tools/pywrekamapd.py"
    def_conf["wrekaUDP_bin"] = "/root/config/code/tools/wrekaUDP.py"
    def_conf["iperf_bin"] = "/usr/local/bin/iperf"
    def_conf["ping_bin"] = "/bin/ping"
    def_conf["wmap2dot_bin"] = "/root/config/code/tools/wmap2dot.awk"
    def_conf["ssh_bin"] = "/usr/bin/ssh"
    def_conf["log_file"] = "wmap.log"
    def_conf["pagrid_out"] = "wmap.out"
    def_conf["pagrid_pdf"] = "wmap.pdf"
    def_conf["pagrid_dot"] = "wmap.dot"
    def_conf["sum_out"] = "wmap.sum"
    def_conf["tmp_dir"] = "/tmp/wmap/"
    def_conf["hostlist"] = "lids16 lids15 lids13 lids09 lids06 lids04 lids02"
    if conf.has_key("tmp_dir"):
        def_conf["hosts_file"] = conf["tmp_dir"] + "hostfile.txt"
    else:
        def_conf["hosts_file"] = ""
    def_conf["mpirun_bin"] = "/usr/local/bin/mpirun"
    def_conf["wrekaMPI_bin"] = "/home/basil/mpi/netMPI/wreka_MPI"
    def_conf["ratio"] = "0 1 0.01"
    def_conf["master_node"] = "0"
    def_conf["cpu_user"] = "basil"
    if conf.has_key("cpu_user") and conf.has_key("hostlist"):
        def_conf["flags_wmapd"] = "--set_list \\\"" + conf["hostlist"] + "\\\" --net_off --set_user " + conf["cpu_user"]
    else:
        def_conf["flags_wmapd"]=""

    if key=="":
        return def_conf
    else:
        return def_conf[key]



def main():
    """ Launch function
    """

    time_start = time.time()
    conf = {}
    
    conf_file = "wmap.conf"
    
    #Parse command line
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hc:", ["help","conf="])
    except getopt.error, msg:
        print msg
        print "For help, use --help"
        sys.exit(2)

    #Process options
    for o, a in opts:
        if o in ("-h", "--help"):
            print __doc__
            sys.exit(0)
        elif o in ("-c", "--conf"):
            conf_file = a
        else:
            print "No such options: " + o
            sys.exit(2)
    
    try:
        handle = file(conf_file, "r")
    except IOError:
        print "ERROR: " + conf_file + " does not exit! Use default values."
#        log.write("File " + conf_file + " does NOT exist. Use default values.")
    else:
        conf = parse_conf(handle, conf)
    
    
    def_conf = default_conf(conf)
    for key in def_conf.keys():
        if key not in conf.keys():
            conf[key] = default_conf(conf, key)

   
    try:
        log = file(conf["log_file"], "w")
    except IOError:
        print "ERROR: " + conf["log_file"] + " impossible to open!\nQuit."
        log.write("File " + conf["log_file"] + " impossible to open. Quit!")
        sys.exit(1)
   
    try:
        os.mkdir(conf["tmp_dir"])
    except:
        log.write("Impossible to create temp dir at " + conf["tmp_dir"] + "\n")
        print "Impossible to create tmp dir at " + conf["tmp_dir"]
    else:
        log.write("Create tmp dir " + conf["tmp_dir"] + " successfully!\n")
 
    try:
        out = file(conf["pagrid_out"], "w")
    except IOError:
        print "ERROR: " + conf["pagrid_out"] + " to open!\nQuit."
        log.write("File " + conf["pagrid_out"] + " impossible to open. Quit!")
        sys.exit(1)
    
    try:
        summ = file(conf["sum_out"], "w")
    except IOError:
        print "ERROR: " + conf["sum_out"] + " to open!\nQuit."
        log.write("File " + conf["sum_out"] + " impossible to open. Quit!")
        sys.exit(1)
    
    user = pwd.getpwuid(os.geteuid())[0] 
    log.write("Logfile for wmap\nGenerated on " + str(datetime.now()) + " by " + user + "\n")
    
    pprint.pprint(conf, log)
    
    mygrid = grid(conf, log)
    mygrid.getInfoNodes()
    mygrid.printSummary()
    
    mypagrid = pagrid(mygrid, log)
    mypagrid.printPagrid(out)
    mypagrid.printGraph()
    #pprint.pprint(mygrid.nodes)

    last = time.time()-time_start
    log.write("Program exited successfully on " + str(datetime.now()) + " (last " + str(last) + "s).\n")

    log.close()
    print "finished OK!"
    
    
if __name__ == "__main__":
    main() 
